package audio

import (
	"encoding/binary"
	"io"
	"os"
)

// WAV Header Constants
const (
	wavHeaderSize = 44
	fmtChunkSize  = 16
	audioFormat   = 1
	bitsPerSample = 16
)

// WriteWavHeader writes the WAV header to the output file
func WriteWavHeader(w io.Writer, numChannels, sampleRate, numSamples int) error {
	byteRate := sampleRate * numChannels * bitsPerSample / 8
	blockAlign := numChannels * bitsPerSample / 8
	dataSize := numSamples * numChannels * bitsPerSample / 8
	wavSize := wavHeaderSize + dataSize - 8

	header := make([]byte, wavHeaderSize)

	// RIFF header
	copy(header[0:], "RIFF")
	binary.LittleEndian.PutUint32(header[4:], uint32(wavSize))
	copy(header[8:], "WAVE")

	// fmt chunk
	copy(header[12:], "fmt ")
	binary.LittleEndian.PutUint32(header[16:], fmtChunkSize)
	binary.LittleEndian.PutUint16(header[20:], audioFormat)
	binary.LittleEndian.PutUint16(header[22:], uint16(numChannels))
	binary.LittleEndian.PutUint32(header[24:], uint32(sampleRate))
	binary.LittleEndian.PutUint32(header[28:], uint32(byteRate))
	binary.LittleEndian.PutUint16(header[32:], uint16(blockAlign))
	binary.LittleEndian.PutUint16(header[34:], bitsPerSample)

	// data chunk
	copy(header[36:], "data")
	binary.LittleEndian.PutUint32(header[40:], uint32(dataSize))

	_, err := w.Write(header)
	return err
}

// ConvertPCMToWAV converts a PCM file to a WAV file
func ConvertPCMToWAV(pcmPath, wavPath string, sampleRate, numChannels int) error {
	pcmFile, err := os.Open(pcmPath)
	if err != nil {
		return err
	}
	defer pcmFile.Close()

	wavFile, err := os.Create(wavPath)
	if err != nil {
		return err
	}
	defer wavFile.Close()

	pcmInfo, err := pcmFile.Stat()
	if err != nil {
		return err
	}

	numSamples := int(pcmInfo.Size()) / (numChannels * bitsPerSample / 8)
	// Write WAV header
	err = WriteWavHeader(wavFile, numChannels, sampleRate, numSamples)
	if err != nil {
		return err
	}
	// Write PCM data to WAV file
	_, err = io.Copy(wavFile, pcmFile)
	return err
}
